# realtime_audio.py
import os
import yaml
import dashscope
import asyncio
import aiofiles
from datetime import datetime
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from dashscope.audio.asr import Recognition, RecognitionCallback, RecognitionResult
from typing import Optional

app = FastAPI()
SAVE_FILE = "recognized_sentences.txt"

# 初始化配置
def load_config():
    """使用与ChatConfig相同的路径解析逻辑"""
    try:
        # 完全复用ChatConfig的路径定位方式
        config_path = os.path.join(
            os.path.dirname(__file__),  # 当前文件所在目录
            "../config.yaml"            # 上级目录的配置文件
        )
        resolved_path = os.path.abspath(config_path)

        # 保持原有错误检查逻辑
        if not os.path.exists(resolved_path):
            raise FileNotFoundError(f"Config file not found at: {resolved_path}")

        with open(resolved_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)

        # 同时检查两种API Key的存在性
        if not config.get("DASHSCOPE_API_KEY"):
            raise KeyError("DASHSCOPE_API_KEY 未在配置文件中找到")
            
        return config["DASHSCOPE_API_KEY"]
        
    except Exception as e:
        # 统一异常处理
        raise RuntimeError(f"配置加载失败: {str(e)}") from e

# 初始化API Key（保持原有错误处理结构）
try:
    dashscope.api_key = load_config()
except Exception as e:
    print(f"初始化失败: {str(e)}")
    exit(1)

class SentenceProcessor:
    """智能句子处理器"""
    def __init__(self):
        self.current_sentence = ""
        self.last_final_sentence = ""
        self.pending_updates = []

    def update(self, text: str, is_end: bool) -> Optional[str]:
        """更新句子状态并返回最终句子"""
        clean_text = text.strip()
        
        if not clean_text:
            return None

        # 更新当前句子
        self.current_sentence = clean_text
        
        # 处理句子结束
        if is_end:
            final_sentence = self._generate_final_sentence()
            if final_sentence and final_sentence != self.last_final_sentence:
                self.last_final_sentence = final_sentence
                return final_sentence
        return None

    def _generate_final_sentence(self) -> str:
        """生成最终句子并重置状态"""
        final_sentence = self.current_sentence
        # 重置状态
        self.current_sentence = ""
        self.pending_updates.clear()
        return final_sentence

class WebSocketCallback(RecognitionCallback):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
        self.processor = SentenceProcessor()
        self.loop = asyncio.get_event_loop()
        self.active = True

    def on_open(self):
        print(f"客户端连接: {self.websocket.client}")

    def on_close(self):
        print(f"连接关闭: {self.websocket.client}")
        self.active = False

    async def _safe_save(self, text: str):
        """带异常处理的保存方法"""
        try:
            async with aiofiles.open(SAVE_FILE, "a", encoding="utf-8") as f:
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                await f.write(f"[{timestamp}] {text}\n")
                print(f"已保存最终句子: {text}")
        except Exception as e:
            print(f"保存失败: {str(e)}")

    def on_event(self, result: RecognitionResult):
        if not self.active:
            return

        sentence = result.get_sentence()
        if 'text' in sentence:
            text = sentence['text']
            is_end = RecognitionResult.is_sentence_end(sentence)
            
            # 处理句子更新
            final_sentence = self.processor.update(text, is_end)
            
            # 保存并发送最终结果
            if final_sentence:
                # 异步发送到客户端
                asyncio.run_coroutine_threadsafe(
                    self.websocket.send_text(final_sentence), self.loop
                )
                
                # 异步保存到文件
                asyncio.run_coroutine_threadsafe(
                    self._safe_save(final_sentence), self.loop
                )

@app.websocket("/ws")
async def handle_audio_stream(websocket: WebSocket):
    await websocket.accept()
    recognizer = None
    
    try:
        callback = WebSocketCallback(websocket)
        recognizer = Recognition(
            model="paraformer-realtime-v2",
            format="pcm",
            sample_rate=16000,
            callback=callback,
            semantic_punctuation_enabled=False,  # 关闭语义断句
            max_sentence_silence=800,  # 800ms静音断句
            disfluency_removal_enabled=True  # 过滤语气词
        )
        recognizer.start()

        # 保持连接直到断开
        while True:
            audio_data = await websocket.receive_bytes()
            recognizer.send_audio_frame(audio_data)
            
    except WebSocketDisconnect:
        print(f"客户端主动断开: {websocket.client}")
    except Exception as e:
        print(f"处理异常: {str(e)}")
    finally:
        try:
            if recognizer:
                recognizer.stop()
        except Exception as e:
            print(f"停止识别器失败: {str(e)}")
        
        if websocket.client_state.name == "CONNECTED":
            await websocket.close(code=1000)

def initialize():
    """初始化服务端"""
    if not os.path.exists(SAVE_FILE):
        open(SAVE_FILE, "w").close()
    print(f"最终句子保存路径: {os.path.abspath(SAVE_FILE)}")

if __name__ == "__main__":
    initialize()
    import uvicorn
    uvicorn.run(
        app, 
        host="0.0.0.0", 
        port=8000,
        ws_ping_timeout=300
    )